package com.img.misc;

import java.awt.image.BufferedImage;
import java.awt.image.WritableRaster;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import javax.imageio.ImageIO;



public class LabelComparison {
	public static class ResultArregator{
		int correct=0;
		int wrong=0;
		public void incrementCorrect(){
			correct++;
		}
		public void incrementWrong(){
			wrong++;
		}
		public double getAccuracy(){
			return ((double)correct*100)/(correct+wrong);
		}
	}
	public static void main(String[] args) throws IOException{
		Map<Character, ResultArregator> class_of_interest=new HashMap<>();
		class_of_interest.put('A', new ResultArregator());
		class_of_interest.put('B', new ResultArregator());
		class_of_interest.put('C', new ResultArregator());
		class_of_interest.put('D', new ResultArregator());
		class_of_interest.put('E', new ResultArregator());
		class_of_interest.put('F', new ResultArregator());
		class_of_interest.put('G', new ResultArregator());
		class_of_interest.put('H', new ResultArregator());
		class_of_interest.put('I', new ResultArregator());
		class_of_interest.put('J', new ResultArregator());
		class_of_interest.put('K', new ResultArregator());
		class_of_interest.put('L', new ResultArregator());
		class_of_interest.put('M', new ResultArregator());
		class_of_interest.put('N', new ResultArregator());
		class_of_interest.put('O', new ResultArregator());
		class_of_interest.put('P', new ResultArregator());
		class_of_interest.put('Q', new ResultArregator());
		class_of_interest.put('R', new ResultArregator());
		class_of_interest.put('S', new ResultArregator());
		class_of_interest.put('T', new ResultArregator());
		class_of_interest.put('U', new ResultArregator());//I am only interested in airplane, so only 'B' is in the set.
		//You may want to add other classes, the classes not in the set will be treated as background.
		//Path ground_truth_folder=Paths.get("D:/VOCdevkit/VOC2012/SegmentationClass");		
		Path ground_truth_folder=Paths.get(args[0]);//Folder for ground truth label image files
		//Path predict_img_folder=Paths.get("D:/ImgPix/hmmGMM4_8/Validate");//folder of images generated by the model
		Path predict_img_folder=Paths.get(args[1]);//folder for predicted label image files
		int[] total_accuracy={0,0,0,0,0,0};
		try(DirectoryStream<Path> dir=Files.newDirectoryStream(predict_img_folder)){
		  for(Path predict_img_path: dir){	
			  
		   try(InputStream ground=Files.newInputStream(ground_truth_folder.resolve(predict_img_path.getFileName()));
			 InputStream result=Files.newInputStream(predict_img_path)){
			 BufferedImage ground_Image=ImageIO.read(ground);
			 BufferedImage predict_Image=ImageIO.read(result);
			 int[] accuracy=compareImg(ground_Image, predict_Image, class_of_interest);
			 //System.out.printf("%s:",predict_img_path.getFileName().toString());
			 //System.out.println(((double)accuracy[0])/accuracy[1]);
			 total_accuracy[0]+=accuracy[0];
			 total_accuracy[1]+=accuracy[1];
			 total_accuracy[2]+=accuracy[2];
			 total_accuracy[3]+=accuracy[3];
			 total_accuracy[4]+=accuracy[4];
			 total_accuracy[5]+=accuracy[5];
		    }
		  }
		}
		System.out.println("Overall:");
		System.out.println(((double)total_accuracy[0])/total_accuracy[1]);
		System.out.println("Background :");
		System.out.println(((double)total_accuracy[2])/total_accuracy[1]);
		System.out.println("Sensitivity :");
		System.out.println(((double)total_accuracy[5])/total_accuracy[3]);
		System.out.println("Specificity :");
		System.out.println(((double)total_accuracy[5])/total_accuracy[4]);
		double sum=0;
		for (Entry<Character, ResultArregator>  es : class_of_interest.entrySet()){
		     System.out.printf("%s : %.2f%n", es.getKey(), es.getValue().getAccuracy());
		     sum+=es.getValue().getAccuracy();
		}
		System.out.printf("[mean] : %.2f%n",  sum/class_of_interest.size());
	}
	
	public static int[] compareImg(BufferedImage ground_truth, BufferedImage predict, Map<Character, ResultArregator>  class_of_interest){ 
		int count=0;
		int correct=0;
		int background=0;
		int foreground=0;
		int predictedforground=0;
		int correctforeground=0;
		WritableRaster ground = ground_truth.getRaster();
		WritableRaster prediction = predict.getRaster();
		for(int y=0; y<ground_truth.getHeight(); y++){
			for(int x=0; x<ground_truth.getWidth(); x++ ){
				char truth_class=getClass(ground, x, y, ground_truth.getType());
				if(truth_class=='Z')
					continue;
				char prediction_class=getClass(prediction, x, y, predict.getType());
				if(!class_of_interest.containsKey(prediction_class))
					prediction_class='A';
				if(!class_of_interest.containsKey(truth_class))
					truth_class='A';
				count++;
				if(truth_class==prediction_class){
					
					correct++;
					if(truth_class!='A'){
						correctforeground++;						
					}
					class_of_interest.get(truth_class).incrementCorrect();
				}else{
					class_of_interest.get(truth_class).incrementWrong();
					class_of_interest.get(prediction_class).incrementWrong();
				}
				if(truth_class=='A')
					background++; 
				else
					foreground++;
				if(prediction_class!='A')
					predictedforground++;
			}
		}
		int[] output={correct, count, background, foreground, predictedforground,correctforeground};
		return output;
	
	}

	private static char getClass(WritableRaster prediction, int x, int y, int type) {
		// TODO Auto-generated method stub
		int[] colorarray=new int[3];
		
		prediction.getPixel(x, y, colorarray);
		if(type==5){
			if (colorarray[0]==0 && colorarray[1]==0 && colorarray[2]==0)
     		    return 'A';
    		else if (colorarray[0]==128 && colorarray[1]==0 && colorarray[2]==0)
    			return 'B';
    		else if (colorarray[0]==0 && colorarray[1]==128 && colorarray[2]==0)
    			return 'C';
    		else if (colorarray[0]==128 && colorarray[1]==128 && colorarray[2]==0)
    			return 'D';
    		else if (colorarray[0]==0 && colorarray[1]==0 && colorarray[2]==128)
    			return 'E';
    		else if (colorarray[0]==128 && colorarray[1]==0 && colorarray[2]==128)
    			return 'F';
    		else if (colorarray[0]==0 && colorarray[1]==128 && colorarray[2]==128)
    			return 'G';
    		else if (colorarray[0]==128 && colorarray[1]==128 && colorarray[2]==128)
    			return 'H';
    		else if (colorarray[0]==64 && colorarray[1]==0 && colorarray[2]==0)
    			return 'I';
    		else if (colorarray[0]==192 && colorarray[1]==0 && colorarray[2]==0)
    			return 'J';
    		else if (colorarray[0]==64 && colorarray[1]==128 && colorarray[2]==0)
    			return 'K';
    		else if (colorarray[0]==192 && colorarray[1]==128 && colorarray[2]==0)
    			return 'L';
    		else if (colorarray[0]==64 && colorarray[1]==0 && colorarray[2]==128)
    			return 'M';
    		else if (colorarray[0]==192 && colorarray[1]==0 && colorarray[2]==128)
    			return 'N';
    		else if (colorarray[0]==64 && colorarray[1]==128 && colorarray[2]==128)
    			return 'O';
    		else if (colorarray[0]==192 && colorarray[1]==128 && colorarray[2]==128)
    			return 'P';
    		else if (colorarray[0]==0 && colorarray[1]==64 && colorarray[2]==0)
    			return 'Q';
    		else if (colorarray[0]==128 && colorarray[1]==64 && colorarray[2]==0)
    			return 'R';
    		else if (colorarray[0]==0 && colorarray[1]==192 && colorarray[2]==0)
    			return 'S';
    		else if (colorarray[0]==128 && colorarray[1]==192 && colorarray[2]==0)
    			return 'T';
    		else if (colorarray[0]==0 && colorarray[1]==64 && colorarray[2]==128)
    			return 'U';
    		else
    			return 'Z';
		}else{
			switch(colorarray[0]){
			case 0: return 'A';
			case 1: return 'B';
			case 2: return 'C';
			case 3: return 'D';
			case 4: return 'E';
			case 5: return 'F';
			case 6: return 'G';
			case 7: return 'H';
			case 8: return 'I';
			case 9: return 'J';
			case 10: return 'K';
			case 11: return 'L';
			case 12: return 'M';
			case 13: return 'N';
			case 14: return 'O';
			case 15: return 'P';
			case 16: return 'Q';
			case 17: return 'R';
			case 18: return 'S';
			case 19: return 'T';
			case 20: return 'U';
			default: return 'Z';
			}
		}
		
		
	}

}
